#!/usr/bin/env python3

import rclpy
from rclpy.node import Node
from rclpy.executors import ExternalShutdownException
from rclpy.qos import qos_profile_sensor_data
from std_msgs.msg import String, Bool
from geometry_msgs.msg import Pose, PoseStamped
from rclpy.action import ActionClient
from moveit_msgs.action import MoveGroup
from moveit_msgs.msg import MotionPlanRequest, Constraints, PositionConstraint, OrientationConstraint
from shape_msgs.msg import SolidPrimitive
import numpy as np
import copy
import time
import threading

# 全局变量
left_pose = None  # 接收到的纸巾目标位姿
yogurt_pose = None  # 接收到的酸奶目标位姿

start_move = False  # 是否开始任务
task_state = "stand_by"  # 任务状态

# 纸巾抓取备选位姿
paper_alternative_poses = [
    {
        "orientation": {
            "x": -0.0321,
            "y": 0.1140,
            "z": 0.9236,
            "w": 0.3645
        }
    },
    {
        "orientation": {
            "x": -0.014,
            "y": 0.332,
            "z": 0.912,
            "w": 0.24
        }
    }
]

# 酸奶抓取备选位姿
yogurt_alternative_poses = [
    {
        "orientation": {
            "x": 0.6855,
            "y": -0.1520,
            "z": 0.6540,
            "w": 0.2812
        }
    },
    {
        "orientation": {
            "x": 0.6402,
            "y": -0.2407,
            "z": 0.6268,
            "w": 0.3731
        }
    }
]

# 固定的释放位置
release_pose = PoseStamped()
release_pose.header.frame_id = "base_link"
release_pose.pose.position.x = 0.5  # 固定释放位置的 x 坐标
release_pose.pose.position.y = -0.55  # 固定释放位置的 y 坐标
release_pose.pose.position.z = 0.85  # 固定释放位置的 z 坐标
release_pose.pose.orientation.x = 0.11236370921869202  # 固定释放位置的姿态
release_pose.pose.orientation.y = 0.15079563957693937
release_pose.pose.orientation.z = 0.7905900885524749
release_pose.pose.orientation.w = 0.5827541366801177


class MoveIt2EEActionClient(Node):
    def __init__(self):
        super().__init__('moveit2_ee_action_client')
        self._action_client = ActionClient(self, MoveGroup, 'move_action')

    def send_goal(self, target_pose):
        """
        发送目标位姿给 MoveIt2，让末端工具移动到目标位置
        :param target_pose: geometry_msgs.msg.Pose 类型的目标位姿
        """
        self.get_logger().info("等待 MoveIt2 服务器...")
        print("target_pose", target_pose)
        self._action_client.wait_for_server()
        # 构造目标请求
        goal_msg = MoveGroup.Goal()
        goal_msg.request.workspace_parameters.header.frame_id = 'base_link'  # 选择世界坐标系
        goal_msg.request.group_name = "right_arm"  # 规划组名称（你的机械臂 MoveIt2 配置）
        goal_msg.request.goal_constraints.append(Constraints())

        # 目标位置约束
        ''''''
        position_constraint = PositionConstraint()
        position_constraint.header.frame_id = 'base_link'  # 目标点的参考坐标系
        position_constraint.link_name = "right_tool_link"  # 末端工具名称
        position_constraint.constraint_region.primitive_poses.append(target_pose)

        # 设置可行区域
        bounding_volume = SolidPrimitive()
        bounding_volume.type = SolidPrimitive.BOX
        bounding_volume.dimensions = [0.0001, 0.0001, 0.0001]  # 允许末端工具偏移的误差范围
        position_constraint.constraint_region.primitives.append(bounding_volume)

        # 目标方向约束
        orientation_constraint = OrientationConstraint()
        orientation_constraint.header.frame_id = 'base_link'
        orientation_constraint.link_name = "right_tool_link"
        orientation_constraint.orientation = target_pose.orientation
        orientation_constraint.absolute_x_axis_tolerance = 0.0001
        orientation_constraint.absolute_y_axis_tolerance = 0.0001
        orientation_constraint.absolute_z_axis_tolerance = 0.0001

        # 加入约束
        goal_msg.request.goal_constraints[0].position_constraints.append(position_constraint)
        goal_msg.request.goal_constraints[0].orientation_constraints.append(orientation_constraint)

        self.get_logger().info("发送目标位姿到 MoveIt2...")
        self._send_goal_future = self._action_client.send_goal_async(goal_msg)
        self._send_goal_future.add_done_callback(self.goal_response_callback)

    def goal_response_callback(self, future):
        goal_handle = future.result()
        if not goal_handle.accepted:
            self.get_logger().error("目标被拒绝")
            return

        self.get_logger().info("MoveIt2 开始执行路径...")
        self._get_result_future = goal_handle.get_result_async()
        self._get_result_future.add_done_callback(self.get_result_callback)

    def get_result_callback(self, future):
        result = future.result().result
        if result.error_code.val == 1:
            self.get_logger().info("末端工具移动成功！")
        else:
            self.get_logger().error(f"移动失败，错误代码：{result.error_code.val}")


class TaskManager(Node):
    def __init__(self):
        super().__init__('task_manager_node')
        self.last_valid_pose = None  # 初始化 last_valid_pose 属性
        self.last_valid_pose_yogurt = None

        self.create_subscription(PoseStamped, '/paper_pose', self.left_pose_callback, 10)
        self.create_subscription(PoseStamped, '/yoger_pose', self.yogurt_pose_callback, 10)
        self.create_subscription(Bool, '/start_move', self.start_move_callback, 10)
        self.hand_control_publisher = self.create_publisher(String, "/hand_control_topic", 10)

        # MoveIt 2 初始化
        self.moveit2_client = MoveIt2EEActionClient()
        self.get_logger().info("Task Manager initialized with MoveIt 2")

    def left_pose_callback(self, msg):
        """
        接收纸巾目标位姿的回调函数。
        """
        if msg.pose.position.x == -1 or msg.pose.position.y == -1 or msg.pose.position.z == -1:
            if self.last_valid_pose:
                self.left_pose = copy.deepcopy(self.last_valid_pose)
            else:
                return
        else:
            if self.last_valid_pose:
                current_position = np.array([msg.pose.position.x, msg.pose.position.y, msg.pose.position.z])
                last_position = np.array(
                    [self.last_valid_pose.pose.position.x, self.last_valid_pose.pose.position.y,
                     self.last_valid_pose.pose.position.z])
                distance = np.linalg.norm(current_position - last_position)

                if distance > 0.10:
                    self.get_logger().info(f"Position change exceeds 1cm ({distance:.4f}m), skipping this pose.")
                    return
                else:
                    self.left_pose = copy.deepcopy(msg)
                    self.last_valid_pose = copy.deepcopy(msg)
            else:
                self.left_pose = copy.deepcopy(msg)
                self.last_valid_pose = copy.deepcopy(msg)
        # self.get_logger().info("++++++++")

    def yogurt_pose_callback(self, msg):
        """
        接收酸奶目标位姿的回调函数。
        """
        if msg.pose.position.x == -1 or msg.pose.position.y == -1 or msg.pose.position.z == -1:
            if self.last_valid_pose_yogurt:
                self.yogurt_pose = copy.deepcopy(self.last_valid_pose_yogurt)
            else:
                self.get_logger().warn('Invalid pose received, and no last valid pose exists.')
                return
        else:
            if self.last_valid_pose_yogurt:
                current_position = np.array([msg.pose.position.x, msg.pose.position.y, msg.pose.position.z])
                last_position = np.array(
                    [self.last_valid_pose_yogurt.pose.position.x, self.last_valid_pose_yogurt.pose.position.y,
                     self.last_valid_pose_yogurt.pose.position.z])
                distance = np.linalg.norm(current_position - last_position)

                if distance > 0.10:
                    self.get_logger().info(f"Position change exceeds 1cm ({distance:.4f}m), skipping this pose.")
                    return
                else:
                    self.yogurt_pose = copy.deepcopy(msg)
                    self.last_valid_pose_yogurt = copy.deepcopy(msg)
            else:
                self.yogurt_pose = copy.deepcopy(msg)
                self.last_valid_pose_yogurt = copy.deepcopy(msg)
        # self.get_logger().info("---------")

    def start_move_callback(self, msg):
        global start_move
        if msg.data:
            self.get_logger().info("Received start_move signal, starting task.")
            start_move = True
        else:
            self.get_logger().info("Received stop_move signal, stopping task.")
            start_move = False

    def move_to_pose(self, pose, link_name, is_yogurt=False):
        print("///")
        self.moveit2_client.send_goal(pose.pose)
        print("+++")
        return True  # Assume success for simplicity

    def move_to_stand_by_pose(self):
        print("**********")
        '''
        stand_by_pose = PoseStamped()
        stand_by_pose.header.frame_id = "base_link"
        stand_by_pose.pose.position.x = 0.4
        stand_by_pose.pose.position.y = -0.55
        stand_by_pose.pose.position.z = 0.80
        stand_by_pose.pose.orientation.x = 0.11236370921869202
        stand_by_pose.pose.orientation.y = 0.15079563957693937
        stand_by_pose.pose.orientation.z = 0.7905900885524749
        stand_by_pose.pose.orientation.w = 0.5827541366801177
        '''
        stand_by_pose = PoseStamped()
        stand_by_pose.header.frame_id = "base_link"
        stand_by_pose.pose.position.x = 0.2
        stand_by_pose.pose.position.y = -0.64
        stand_by_pose.pose.position.z = 0.9
        # stand_by_pose.pose.orientation.x = 0.11236370921869202
        # stand_by_pose.pose.orientation.y = 0.15079563957693937
        # stand_by_pose.pose.orientation.z = 0.7905900885524749
        stand_by_pose.pose.orientation.w = 1.0
        return self.move_to_pose(stand_by_pose, "right_tool_link")

    def adjust_pose_for_paper(self, pose):
        adjusted_pose = copy.deepcopy(pose)
        adjusted_pose.pose.position.x += 0.02
        adjusted_pose.pose.position.y += -0.10
        adjusted_pose.pose.position.z += 0.03
        adjusted_pose.pose.orientation.x = 0.0927
        adjusted_pose.pose.orientation.y = 0.2553
        adjusted_pose.pose.orientation.z = 0.9106
        adjusted_pose.pose.orientation.w = 0.3113
        return adjusted_pose

    def adjust_pose_for_yogurt(self, pose):
        adjusted_pose = copy.deepcopy(pose)
        adjusted_pose.pose.position.x += 0.05
        adjusted_pose.pose.position.y += -0.05
        adjusted_pose.pose.position.z += -0.06
        adjusted_pose.pose.orientation.x = 0.6878
        adjusted_pose.pose.orientation.y = -0.0545
        adjusted_pose.pose.orientation.z = 0.6945
        adjusted_pose.pose.orientation.w = 0.2037
        return adjusted_pose

    def task_manager(self):
        self.get_logger().info("task_manager_start_+++++++")
        global task_state, start_move, left_pose, yogurt_pose, hand_control_publisher, release_pose
        self.get_logger().info("task_state_+++++++")
        print("task_state", task_state)
        rate = self.create_rate(10)  # 10Hz
        self.get_logger().info("createe_+++++++")
        while rclpy.ok():
            self.get_logger().info("rclpy_+++++++")

            if not start_move:
                rate.sleep()
                continue


            if task_state == "stand_by":
                self.get_logger().info("Task state: stand_by")

                if self.move_to_stand_by_pose():
                    self.get_logger().info("Arm is now in stand_by pose.")
                else:
                    self.get_logger().warn("Failed to move to stand_by pose, retrying...")
                    continue

                if left_pose is not None:
                    task_state = "move_above_paper"
                else:
                    self.get_logger().info("Waiting for target pose...")

            elif task_state == "move_above_paper":
                self.get_logger().info("Task state: move_above_paper")
                adjusted_pose = self.adjust_pose_for_paper(left_pose)
                above_pose = copy.deepcopy(adjusted_pose)
                above_pose.pose.position.z += 0.10
                if self.move_to_pose(above_pose, "right_tool_link"):
                    task_state = "move_to_paper"
                    time.sleep(1)
                else:
                    self.get_logger().warn("Failed to move above paper, retrying...")

            elif task_state == "move_to_paper":
                self.get_logger().info("Task state: move_to_paper")
                adjusted_pose = self.adjust_pose_for_paper(left_pose)
                if self.move_to_pose(adjusted_pose, "right_tool_link"):
                    task_state = "grasp"
                    time.sleep(1)
                else:
                    self.get_logger().warn("Failed to move to paper pose, retrying...")

            elif task_state == "grasp":
                self.get_logger().info("Task state: grasp")
                msg_hand_control = String()
                msg_hand_control.data = "paper_grasp"
                self.hand_control_publisher.publish(msg_hand_control)
                self.get_logger().info("Sent hand close command for paper.")
                time.sleep(2)
                task_state = "move_up_after_grasp"
                time.sleep(1)

            elif task_state == "move_up_after_grasp":
                self.get_logger().info("Task state: move_up_after_grasp")
                current_pose = self.moveit2_client.get_current_pose().pose
                target_pose = copy.deepcopy(current_pose)
                target_pose.position.z += 0.10
                target_pose_stamped = PoseStamped()
                target_pose_stamped.header.frame_id = "base_link"
                target_pose_stamped.pose = target_pose
                if self.move_to_pose(target_pose_stamped, "right_tool_link"):
                    task_state = "move_to_release"
                    time.sleep(1)
                else:
                    self.get_logger().warn("Failed to move up after grasp, retrying...")

            elif task_state == "move_to_release":
                self.get_logger().info("Task state: move_to_release")
                if self.move_to_pose(release_pose, "right_tool_link"):
                    task_state = "release"
                    time.sleep(1)
                else:
                    self.get_logger().warn("Failed to move to release pose, retrying...")

            elif task_state == "release":
                self.get_logger().info("Task state: release")
                msg_hand_control = String()
                msg_hand_control.data = "release"
                self.hand_control_publisher.publish(msg_hand_control)
                self.get_logger().info("Sent hand open command for paper.")
                time.sleep(2)
                task_state = "return_home_after_paper"
                time.sleep(1)

            elif task_state == "return_home_after_paper":
                self.get_logger().info("Task state: return_home_after_paper")
                if self.move_to_stand_by_pose():
                    self.get_logger().info("Returned to stand_by pose after paper task.")
                    task_state = "move_above_yogurt"
                    time.sleep(1)
                else:
                    self.get_logger().warn("Failed to return to stand_by pose after paper task, retrying...")

            elif task_state == "move_above_yogurt":
                self.get_logger().info("Task state: move_above_yogurt")
                if yogurt_pose is None:
                    self.get_logger().info("Waiting for yogurt pose...")
                    rate.sleep()
                    continue
                adjusted_yogurt_pose = self.adjust_pose_for_yogurt(yogurt_pose)
                above_pose = copy.deepcopy(adjusted_yogurt_pose)
                above_pose.pose.position.x += 0.01
                above_pose.pose.position.y += -0.10
                above_pose.pose.position.z += 0.16
                if self.move_to_pose(above_pose, "right_tool_link", is_yogurt=True):
                    task_state = "move_to_yogurt"
                    time.sleep(1)
                else:
                    self.get_logger().warn("Failed to move above yogurt, retrying...")

            elif task_state == "move_to_yogurt":
                self.get_logger().info("Task state: move_to_yogurt")
                adjusted_yogurt_pose = self.adjust_pose_for_yogurt(yogurt_pose)
                if self.move_to_pose(adjusted_yogurt_pose, "right_tool_link", is_yogurt=True):
                    task_state = "grasp_yogurt"
                    time.sleep(1)
                else:
                    self.get_logger().warn("Failed to move to yogurt pose, retrying...")

            elif task_state == "grasp_yogurt":
                self.get_logger().info("Task state: grasp_yogurt")
                msg_hand_control = String()
                msg_hand_control.data = "yogurt_grasp"
                self.hand_control_publisher.publish(msg_hand_control)
                self.get_logger().info("Sent hand close command for yogurt.")
                time.sleep(2)
                task_state = "move_up_after_grasp_yogurt"
                time.sleep(1)

            elif task_state == "move_up_after_grasp_yogurt":
                self.get_logger().info("Task state: move_up_after_grasp_yogurt")
                current_pose = self.moveit2_client.get_current_pose().pose
                target_pose = copy.deepcopy(current_pose)
                target_pose.position.z += 0.15
                target_pose_stamped = PoseStamped()
                target_pose_stamped.header.frame_id = "base_link"
                target_pose_stamped.pose = target_pose
                if self.move_to_pose(target_pose_stamped, "right_tool_link", is_yogurt=True):
                    task_state = "move_to_yogurt_release"
                    time.sleep(1)
                else:
                    self.get_logger().warn("Failed to move up after yogurt grasp, retrying...")

            elif task_state == "move_to_yogurt_release":
                self.get_logger().info("Task state: move_to_yogurt_release")
                if self.move_to_pose(release_pose, "right_tool_link", is_yogurt=True):
                    task_state = "release_yogurt"
                    time.sleep(1)
                else:
                    self.get_logger().warn("Failed to move to yogurt release pose, retrying...")

            elif task_state == "release_yogurt":
                self.get_logger().info("Task state: release_yogurt")
                msg_hand_control = String()
                msg_hand_control.data = "release"
                self.hand_control_publisher.publish(msg_hand_control)
                self.get_logger().info("Sent hand open command for yogurt.")
                time.sleep(2)
                task_state = "return_home_after_yogurt"
                time.sleep(1)

            elif task_state == "return_home_after_yogurt":
                self.get_logger().info("Task state: return_home_after_yogurt")
                if self.move_to_stand_by_pose():
                    self.get_logger().info("Task completed, returning to stand_by state.")
                    task_state = "stand_by"
                    start_move = False
                    global last_valid_pose, last_valid_pose_yogurt
                    last_valid_pose = None
                    last_valid_pose_yogurt = None
                    time.sleep(1)
                else:
                    self.get_logger().warn("Failed to return to stand_by pose after yogurt task, retrying...")

            rate.sleep()


def main(args=None):
    # 初始化 ROS 2 系统
    rclpy.init(args=args)

    # 创建 TaskManager 节点
    task_manager = TaskManager()

    # 创建 MoveIt2 客户端
    moveit2_client = MoveIt2EEActionClient()

    # 进入 ROS 2 执行循环
    try:
        rclpy.spin(task_manager)
    except KeyboardInterrupt:
        pass
    finally:
        # 关闭 ROS 2 系统
        task_manager.destroy_node()
        rclpy.shutdown()

if __name__ == '__main__':
    main()
